import os
import numpy as np
from pathlib import Path
from collections import defaultdict


def convert_npy_to_bin_simple(src_dir, target_dir):
    """
    将npy文件转换为bin文件

    源结构: src_dir/{operator}/{framework}/*.npy
    目标结构: target_dir/{operator}/round_{i}/input_0.bin

    忽略framework，将所有文件按顺序转换为round
    """
    src_path = Path(src_dir)
    target_path = Path(target_dir)

    # 遍历所有算子
    for op_dir in src_path.iterdir():
        if not op_dir.is_dir():
            continue

        op_name = op_dir.name
        print(f"Processing operator: {op_name}")

        # 收集所有framework下的npy文件
        all_files = []
        for framework_dir in op_dir.iterdir():
            if not framework_dir.is_dir():
                continue

            for npy_file in framework_dir.glob('*.npy'):
                all_files.append(npy_file)

        # 按文件名排序（包含时间戳，保证顺序）
        all_files.sort(key=lambda x: x.name)

        print(f"  Found {len(all_files)} files")

        # 转换每个文件为一个round
        for round_idx, npy_file in enumerate(all_files):
            # 读取npy文件
            data = np.load(npy_file)

            # 确保是float32格式
            if data.dtype != np.float32:
                print(f"  Warning: Converting {npy_file.name} from {data.dtype} to float32")
                data = data.astype(np.float32)

            # 创建目标目录
            round_dir = target_path / op_name / f"round_{round_idx}"
            round_dir.mkdir(parents=True, exist_ok=True)

            # 保存为bin文件
            bin_filepath = round_dir / "input_0.bin"
            data.tofile(bin_filepath)

            if round_idx % 50 == 0:
                print(f"    Processed round_{round_idx}")

        print(f"  ✓ Converted {len(all_files)} files to rounds")
        print(f"    Output: {target_path / op_name}")


def convert_with_counter_grouping(src_dir, target_dir):
    """
    按counter分组转换（如果需要保持原始counter逻辑）

    这个版本会将相同counter的文件放在同一个round
    """
    src_path = Path(src_dir)
    target_path = Path(target_dir)

    # 遍历所有算子
    for op_dir in src_path.iterdir():
        if not op_dir.is_dir():
            continue

        op_name = op_dir.name
        print(f"Processing operator: {op_name}")

        # 收集所有文件并按counter分组
        counter_groups = defaultdict(list)

        for framework_dir in op_dir.iterdir():
            if not framework_dir.is_dir():
                continue

            framework_name = framework_dir.name

            for npy_file in framework_dir.glob('*.npy'):
                # 从文件名提取counter
                # 格式: MRE_random_00001_20231102_153000.npy
                parts = npy_file.stem.split('_')

                # 找到counter部分（5位数字）
                counter = None
                for part in parts:
                    if part.isdigit() and len(part) == 5:
                        counter = int(part)
                        break

                if counter is not None:
                    # 存储 (文件路径, framework, 时间戳)
                    timestamp = parts[-2] + parts[-1] if len(parts) >= 2 else ""
                    counter_groups[counter].append((npy_file, framework_name, timestamp))

        print(f"  Found {len(counter_groups)} unique counters")

        # 处理每个counter组
        for counter in sorted(counter_groups.keys()):
            files = counter_groups[counter]

            # 如果同一个counter有多个文件，选择第一个（或可以选择特定framework）
            # 这里按时间戳排序，选择最早的
            files.sort(key=lambda x: x[2])
            npy_file, framework, timestamp = files[0]

            # 读取并转换
            data = np.load(npy_file)

            # 确保是float32
            if data.dtype != np.float32:
                data = data.astype(np.float32)

            # 创建round目录
            round_dir = target_path / op_name / f"round_{counter}"
            round_dir.mkdir(parents=True, exist_ok=True)

            # 保存
            bin_filepath = round_dir / "input_0.bin"
            data.tofile(bin_filepath)

            if counter % 50 == 0:
                print(f"    Processed round_{counter} from {framework}")

        print(f"  ✓ Converted to {len(counter_groups)} rounds")


def verify_conversion(src_file, bin_file):
    """验证转换的正确性"""
    # 读取原始npy
    original = np.load(src_file)

    # 读取bin文件
    converted = np.fromfile(bin_file, dtype=np.float32)
    converted = converted.reshape(original.shape)

    # 比较
    if np.allclose(original.astype(np.float32), converted):
        return True, "Match"
    else:
        max_diff = np.max(np.abs(original.astype(np.float32) - converted))
        return False, f"Mismatch, max diff: {max_diff}"


def batch_convert_and_verify(src_dir, target_dir, verify=True):
    """转换并验证"""
    print("=" * 80)
    print("NPY to BIN Converter for C++ Loading")
    print("=" * 80)
    print(f"Source: {src_dir}")
    print(f"Target: {target_dir}")
    print("=" * 80)
    print()

    # 执行转换
    convert_npy_to_bin_simple(src_dir, target_dir)

    # 验证（抽样检查）
    if verify:
        print("\n" + "=" * 80)
        print("Verification (sampling)")
        print("=" * 80)

        src_path = Path(src_dir)
        target_path = Path(target_dir)

        for op_dir in src_path.iterdir():
            if not op_dir.is_dir():
                continue

            op_name = op_dir.name

            # 随机选择几个文件验证
            all_npy = []
            for framework_dir in op_dir.iterdir():
                if framework_dir.is_dir():
                    all_npy.extend(list(framework_dir.glob('*.npy')))

            all_npy.sort(key=lambda x: x.name)

            # 验证前5个
            for i, npy_file in enumerate(all_npy[:5]):
                round_dir = target_path / op_name / f"round_{i}"
                bin_file = round_dir / "input_0.bin"

                if bin_file.exists():
                    success, msg = verify_conversion(npy_file, bin_file)
                    status = "✓" if success else "✗"
                    print(f"{status} {op_name}/round_{i}: {msg}")

    print("\n" + "=" * 80)
    print("Conversion completed!")
    print("=" * 80)


if __name__ == "__main__":
    # 配置路径
    SOURCE_DIR = "problem_tensors"  # 源目录
    TARGET_DIR = "target_dir"  # 目标目录

    # 方式1: 简单按顺序转换（推荐）
    batch_convert_and_verify(SOURCE_DIR, TARGET_DIR, verify=True)

    # 方式2: 按counter分组（如果需要保持counter逻辑）
    # convert_with_counter_grouping(SOURCE_DIR, TARGET_DIR)
